# -*- coding: utf-8 -*-

# -------------------------------
# Author    : linlingwei
# Email     : 951035650@qq.com
# File Name : rpIkhandle.py
# Date      : 2019/12/10 
# IDE       : PyCharm
# Version   : 1.0.1
# -------------------------------

import maya.cmds as mc
import maya.mel as mel

import configs.constants.attribute_constant as ac
import configs.constants.nodeType_constant as nc
import configs.constants.suffix_constant as sc
import configs.constants.body_constant as bc
import configs.constants.global_constant as gc

import scripts.commonCmds as con
import scripts.curveCmds as crc

import rig.ikCore as ikCore
reload(ikCore)

class RpIkhandle(ikCore.IkCore):
    def __init__(self,lname="test"):
        super(RpIkhandle,self).__init__("ikRPsolver",lname=lname)
        self.polarObj = None
        
    def create(self,sjot,ejot,polar=None):
        self._init()
        self.ikhandle(sjot,ejot)
        if polar:
            self.polar(polar)
        return
    
    def auto_polar(self,pro=2):
        
        tmploc = mc.spaceLocator(name="tmp")[0]
        mc.delete(mc.pointConstraint(self.jotlist[0], self.jotlist[-1], tmploc, mo=False))
        #aimConstraint -offset 0 0 0 -weight 1 -aimVector 1 0 0 -upVector 0 1 0 -worldUpType "object" -worldUpObject locator2;

        mc.delete(mc.aimConstraint(self.jotlist[-1], tmploc, mo=False,
                                   aimVector=[1 ,0, 0], upVector =[0,0,-1], worldUpType= "object",
                                   worldUpObject=self.jotlist[1]))
        
        grp = con.emptyGroup(tmploc)
        dis = con.point_distance(self.jotlist[0], self.jotlist[-1])
        mc.setAttr(tmploc+".tz",-dis*pro)
        
        loc = con.objToLocator(tmploc,con.lname(self.lname,bc.gPolar),parent=False)
        mc.delete(grp[0])
        
        self.polar(loc)
        return loc
    
    def polar_follow(self):
        
        star_loc = con.objToLocator(self.jotlist[0],name=con.lname(self.lname,"star_polar_flw",))
        end_loc = con.objToLocator(self.jotlist[-1],name=con.lname(self.lname,"end_polar_flw"))
        flwloc = mc.spaceLocator(name=con.lname(self.lname,"polar_follow"))[0]
        mc.parent(flwloc,self.localGrp)
        flw_grp = con.emptyGroup(flwloc)
        mc.pointConstraint(star_loc, end_loc,flw_grp[1], mo=False)
        #aimConstraint -offset 0 0 0 -weight 1 -aimVector 1 0 0 -upVector 0 1 0 -worldUpType "object" -worldUpObject locator2;

        mc.aimConstraint(end_loc, flw_grp[2], mo=False,
                                   aimVector=[1 ,0, 0], upVector =[0,0,-1], worldUpType= "none",
                                   # worldUpObject=self.jotlist[1]
                         )
        mc.parent(star_loc,self.jotRoot)
        mc.parent(end_loc,self.ikhandle_loc)
        
        mc.hide(star_loc,end_loc,flwloc)
        
        return flwloc
    
    def polar(self,pobj):
        mc.poleVectorConstraint(pobj,self._ikhandle)
        mc.parent(pobj,self.localGrp)
        np = mc.rename(pobj,con.lname(self.lname,bc.gPolar))
        self.polarObj = np
        
        con.createAnn(self.jotlist[1],np)
        return np

    def initJoint(self, axial='x'):
        # ikhandle = self._ikhandle
        if 'init' in mc.listAttr(self.localBridge):
            return
        mc.addAttr(self.localBridge, ln='init', at='bool', dv=True, k=False)
        con.lockHideAttr(self.localBridge, ['init'], lock=True, hide=True)
        
        if ac.global_scale not in mc.listAttr(self.localBridge):
            mc.addAttr(self.localBridge, ln=ac.global_scale, at='float', dv=1, k=True)
        mc.addAttr(self.localBridge, ln=ac.local_scale, at='float', dv=1, min=0.01, k=True)
    
        jotList = self.jotlist
    
        attr = ac.lock_polar
        if attr not in mc.listAttr(self.localBridge):
            mc.addAttr(self.localBridge, ln=attr, at='float', dv=0,k=True,min=0,max=1)
            mc.setAttr('{}.{}'.format(self.localBridge, attr), 0)
    
        Length = ac.local_length
        for i, jnt in enumerate(jotList[:-1]):
            orgJntLen = mc.getAttr(jotList[i+1] + '.t%s' % axial)
            jntLen = abs(orgJntLen)
        
            attr = Length
            if attr not in mc.listAttr(jnt):
                mc.addAttr(jnt, ln=attr, at='float', dv=jntLen)
                mc.setAttr('{}.{}'.format(jnt, attr), jntLen, l=True)
            attr = ac.offset
            if attr not in mc.listAttr(jnt):
                mc.addAttr(jnt, ln=attr, at='float', dv=0)
                mc.setAttr('{}.{}'.format(jnt, attr), 0, k=True, )
            attr = ac.stretch
            if attr not in mc.listAttr(jnt):
                mc.addAttr(jnt, ln=attr, at='float', dv=0)
                mc.setAttr('{}.{}'.format(jnt, attr), 0, k=True, )
            attr = ac.lock_polar
            if attr not in mc.listAttr(jnt):
                mc.addAttr(jnt, ln=attr, at='float', dv=jntLen)
                mc.setAttr('{}.{}'.format(jnt, attr), jntLen, k=True, )
        
            attr = ac.local_scale
            if attr not in mc.listAttr(jnt):
                mc.addAttr(jnt, ln=attr, at='float', dv=1, min=0.01, k=True)
                lsMd = mc.createNode(nc.tMultiplyDivide, name=con.lname(jnt, ac.local_scale, sc.sMd))
                mc.connectAttr(self.localBridge + "." + ac.local_scale, lsMd + ".input1X")
                mc.connectAttr(jnt + "." + ac.local_scale, lsMd + ".input2X")
            
                mc.connectAttr(lsMd + ".outputX", jnt + ".sx")
                mc.connectAttr(lsMd + ".outputX", jnt + ".sy")
                mc.connectAttr(lsMd + ".outputX", jnt + ".sz")
        
            jntLenBta = mc.createNode('blendTwoAttr', name=con.lname(jnt, Length, sc.sBta))
            mc.connectAttr(jnt + ".{}".format(ac.lock_polar), jntLenBta + '.input[1]')
            negMd = mc.createNode(nc.tMultiplyDivide, name=con.lname(jnt, 't%s' % axial, "Abs", sc.sMd))
            mc.connectAttr(jntLenBta + '.output', negMd + '.input1X')
            mc.connectAttr(negMd + '.outputX', jotList[i+1] + '.t%s' % axial)
        
            # localMd = mc.createNode(tMultiplyDivide,name=cCore.lname(jnt, 't%s' % axial, aLocalScale, sMd))
            # mc.connectAttr(negMd + '.outputX',localMd+".input1X")
            # mc.connectAttr(jnt+"."+aLocalScale,localMd+".input2X")
            # mc.setAttr(localMd+".op",2)
            # mc.connectAttr(localMd + '.outputX', jnt + '.t%s' % axial)
        
            if orgJntLen < 0:
                mc.setAttr(negMd + ".input2X", -1)
            else:
                mc.setAttr(negMd + ".input2X", 1)
        
            mc.connectAttr(self.localBridge + ".{}".format(ac.lock_polar), jntLenBta + ".attributesBlender")
        
            jntLenPma = mc.createNode('plusMinusAverage',
                                      name=con.lname(jnt, Length, sc.sPma))
            mc.connectAttr(jntLenPma + '.output1D', jntLenBta + '.input[0]')
        
            mc.connectAttr(jnt + '.{}'.format(Length), jntLenPma + '.input1D[0]')
        
            mc.setAttr(jntLenPma + '.operation', 1)
        
            mc.connectAttr(jnt + '.'+ac.offset, jntLenPma + '.input1D[1]')
            mc.setAttr(jntLenPma + '.operation', 1)
        
            mc.connectAttr(jnt + '.stretch', jntLenPma + '.input1D[2]')
            mc.setAttr(jntLenPma + '.operation', 1)
    
        return
    def stretch(self):
        jotlist = self.jotlist
        # jotgrp = con.emptyGroup(jotlist[0],["Zero"])[0]
        self.initJoint()

        if self.localBridge is not None:
            bre = self.localBridge
        else:
            bre = self._ikhandle

        mc.addAttr(bre, ln=ac.stretch_switch, at='long', dv=1, min=0, max=1, k=True)
        
        len_cst_con = mc.createNode('condition', name=con.lname(self.lname, 'contrast', sc.sCon))
        mc.setAttr(len_cst_con + ".operation", 2)

        pro_md = mc.createNode('multiplyDivide', name=con.lname(self.lname, 'pro', sc.sMd))
        starLoc = con.objToLocator(jotlist[0], con.lname(self.lname, ac.stretch,"star"), parent=False)
        ikhandleLoc = con.objToLocator(self._ikhandle, con.lname(self.lname, ac.stretch,"end"), parent=False)
        
        mc.parent(starLoc,self.jotRoot)
        mc.parent(ikhandleLoc,self.ikhandle_loc)
        
        dyn_dbn = mc.createNode('distanceBetween', name=con.lname(self.lname, 'TotalLength', sc.sDbn))
        mc.connectAttr(starLoc + ".worldMatrix[0]", dyn_dbn + ".inMatrix1")
        mc.connectAttr(ikhandleLoc + ".worldMatrix[0]", dyn_dbn + ".inMatrix2")

        lgMd = mc.createNode(nc.tMultiplyDivide, name=con.lname(self.lname, ac.local_scale, sc.sMd))
        mc.connectAttr(dyn_dbn + ".distance", lgMd + ".input1X")
        mc.connectAttr(self.localBridge + "." + ac.local_scale, lgMd + ".input2X")
        mc.setAttr(lgMd + ".operation", 2)

        gloMd = mc.createNode('multiplyDivide', name=con.lname(self.lname, 'gloPro', sc.sMd))
        # mc.connectAttr(dyn_dbn + ".distance", gloMd + ".input1X")
        mc.connectAttr(lgMd + ".outputX", gloMd + ".input1X")
        mc.connectAttr(self.localBridge + "." + ac.global_scale, gloMd + ".input2X")
        mc.setAttr(gloMd + ".operation", 2)

        org_pma = mc.createNode('plusMinusAverage', name=con.lname(self.lname, 'org', sc.sPma))
        for i, jot in enumerate(jotlist[:-1]):
            mc.connectAttr(jot + "." + ac.local_length, org_pma + ".input1D[%s]" % i)

        # len_cst_con = mc.createNode('condition', name=gel.name(ikhandle, make='contrast', suffix=cot.con))
        mc.connectAttr(gloMd + ".outputX", len_cst_con + ".firstTerm")
        mc.connectAttr(org_pma + ".output1D", len_cst_con + ".secondTerm")
        mc.connectAttr(gloMd + ".outputX", len_cst_con + ".colorIfTrueR")
        mc.connectAttr(org_pma + ".output1D", len_cst_con + ".colorIfFalseR")

        mc.connectAttr(org_pma + ".output1D", pro_md + ".input2X")
        mc.connectAttr(len_cst_con + ".outColorR", pro_md + ".input1X")
        mc.setAttr(pro_md + ".operation", 2)

        pro_pma = mc.createNode('plusMinusAverage', name=con.lname(self.lname, "pro", sc.sPma))
        mc.setAttr(pro_pma + ".input1D[1]", 1)
        mc.connectAttr(pro_md + ".outputX", pro_pma + ".input1D[0]")
        mc.setAttr(pro_pma + ".operation", 2)

        stretch_con = mc.createNode('condition', name=con.lname(self.lname, "switch", sc.sCon))
        mc.connectAttr(bre + "." + ac.stretch_switch, stretch_con + ".firstTerm")
        mc.setAttr(stretch_con + ".secondTerm", 1)
        mc.setAttr(stretch_con + ".operation", 0)
        mc.connectAttr(pro_pma + ".output1D", stretch_con + ".colorIfTrueR")
        mc.setAttr(stretch_con + ".colorIfFalseR", 0)

        mc.addAttr(bre, ln=ac.stretch_ratio, at='float', dv=0, k=True)
        mc.connectAttr(stretch_con + ".outColorR", bre + "." + ac.stretch_ratio)

        for jot in jotlist[:-1]:
            proAdd_md = mc.createNode('multiplyDivide', name=con.lname(jot, 'stretch', sc.sMd))
            mc.connectAttr(stretch_con + ".outColorR", proAdd_md + ".input1X")
            mc.connectAttr(jot + "." + ac.local_length, proAdd_md + ".input2X")
            mc.setAttr(proAdd_md + ".operation", 1)
            mc.connectAttr(proAdd_md + ".outputX", jot + "." + ac.stretch)

        return starLoc, ikhandleLoc

    def lockPolar(self):
        if not self.polarObj:
            return
        self.initJoint()
        jointList = self.jotlist
        polar = self.polarObj
        upLock = con.objToLocator(jointList[0], con.lname(self.lname, 'upLock'), parent=False)
        polarLock = con.objToLocator(polar, con.lname(self.lname, ac.lock_polar), parent=False)
        dnLock = con.objToLocator(jointList[2], con.lname(self.lname, 'dnLock'), parent=False)
        
        mc.parent(upLock,self.jotRoot)
        mc.parent(polarLock,polar)
        mc.parent(dnLock,self.ikhandle_loc)
        
        updyn_dbn = mc.createNode('distanceBetween', name=con.lname(self.lname, 'upLockKnee', sc.sDbn))
        mc.connectAttr(upLock + ".worldMatrix[0]", updyn_dbn + ".inMatrix1")
        mc.connectAttr(polarLock + ".worldMatrix[0]", updyn_dbn + ".inMatrix2")
    
        dndyn_dbn = mc.createNode('distanceBetween', name=con.lname(self.lname, 'dnLockKnee', sc.sDbn))
        mc.connectAttr(dnLock + ".worldMatrix[0]", dndyn_dbn + ".inMatrix1")
        mc.connectAttr(polarLock + ".worldMatrix[0]", dndyn_dbn + ".inMatrix2")
    
        gloMd = mc.createNode('multiplyDivide', name=con.lname(self.lname, 'gloLockKnee', sc.sMd))
        mc.connectAttr(updyn_dbn + ".distance", gloMd + ".input1X")
        mc.connectAttr(dndyn_dbn + ".distance", gloMd + ".input1Y")
    
        mc.connectAttr(self.localBridge + "." + ac.global_scale, gloMd + ".input2X")
        mc.connectAttr(self.localBridge + "." + ac.global_scale, gloMd + ".input2Y")
    
        mc.setAttr(gloMd + ".operation", 2)
        mc.connectAttr(gloMd + ".outputX", jointList[0] + "." + ac.lock_polar)
        mc.connectAttr(gloMd + ".outputY", jointList[1] + "." + ac.lock_polar)
    
        # mc.parent(polarLock, polar)
    
        # mc.addAttr(self.localBridge, ln=ac.lock_polar, at='long', dv=0, min=0, max=1, k=True)
        # mc.connectAttr(self.localBridge + '.' + ac.lock_polar, ikhandle + '.' + ac.lock_polar)
        return upLock, dnLock, polarLock
    

if __name__ == '__main__':
    pass